import sys
from utils import *
from tqdm import tqdm as tqdm
import time
import matplotlib.pyplot as plt

EPOCH = 60
Ks = [3, 5, 7]
Vs = []

# 加载数据
trainDataset, trainLabels = loadDataSet("./ORL", "train")
print("训练数据加载完毕!")
testDataset, testLabels = loadDataSet("./ORL", "test")
print("测试数据加载完毕!")

# 主程序
for K in Ks:
    accuracy = 0
    wrongSamples = None

    with tqdm(range(EPOCH), file=sys.stdout) as bar:
        for epoch in bar:
            centroids, clusterStandard = train(trainDataset, trainLabels, K)
            t, ws = test(centroids, clusterStandard, testDataset, testLabels, K)
            if accuracy < t:
                accuracy = t
                wrongSamples = ws
            bar.set_description("K值:{:d} 准确率:{:.6f}".format(K, accuracy))
    time.sleep(0.1)
    bar.close()
    Vs.append(accuracy)
    print("错误样本:", wrongSamples)

plt.xlabel("K")
plt.ylabel("Accuracy")
plt.plot(Ks, Vs)
for a, b in zip(Ks, Vs):
    plt.text(a, b, b)
plt.show()
